#!/bin/bash
# Copyright 2024 CHINA MERCHANTS BANK CO., LTD.
# Copyright 2024 Huawei Technologies Co., Ltd
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================

if [ $# != 4 ]
then
  echo "Usage Help: bash dry_run.sh [CONFIG_PATH] [RANK_SIZE] [TRAIN_DATASET_DIR] [SEQ_LEN]"
  exit 1
fi

#!/bin/bash
YAML_FILE=$1
RANK_SIZE=$2
TRAIN_DATASET_DIR=$3
SEQ_LEN=$4
ROOT_PATH=`pwd`

export MS_SIMULATION_LEVEL=1
export RANK_SIZE=$RANK_SIZE

rm -rf output_dryrun
mkdir output_dryrun
for((i=0; i<$RANK_SIZE; i++))
do
    export DEVICE_ID=$i
    export RANK_ID=$i
    echo "start training for rank $RANK_ID, device $DEVICE_ID"
    # 需要正确指定 run_mindformer.py 路径
    python ./run_train.py \
            --config $ROOT_PATH/$1 \
            --only_save_strategy True \
            --train_dataset_dir $3 \
            --seq_length $4 \
            --output_dir ./output_dryrun \
            --auto_trans_ckpt False &> ./output_dryrun/rank_$RANK_ID.log &
done
